#!/bin/bash
# Huawei CDC NCM Module Build and Installation Script
# This script automates the process of building and installing the huawei_cdc_ncm module
# for the NVIDIA Jetson Developer Kit.
#
# Usage:
#   ./make_huawei_cdc_ncm_module.sh [[-d directory ] | [-h]]
#
# Options:
#   -d | --directory  Specify the kernel source directory (default: /usr/src)
#   -h | --help       Display this help message

# Color functions for consistent output
red=$(tput setaf 1)
green=$(tput setaf 2)
yellow=$(tput setaf 3)
reset=$(tput sgr0)

echo_error() {
    echo "${red}Error: $1${reset}" >&2
}

echo_info() {
    echo "$1"
}

echo_warning() {
    echo "${yellow}Warning: $1${reset}"
}

echo_success() {
    echo "${green}$1${reset}"
}

# Usage function for help message
function usage {
    echo "usage: ./make_huawei_cdc_ncm_module.sh [[-d directory ]  | [-h]]"
    echo "-d | --directory  Directory path to parent of kernel source (default: /usr/src)"
    echo "-h | --help       Show this help message"
}

# Default source directory
SOURCE_TARGET="/usr/src"

# Parse command-line arguments
while [ "$1" != "" ]; do
    case $1 in
        -d | --directory ) shift
            SOURCE_TARGET=$1
            ;;
        -h | --help )
            usage
            exit 0
            ;;
        * )
            usage
            exit 1
    esac
    shift
done

# Ensure SOURCE_TARGET ends with a slash
[[ "${SOURCE_TARGET}" != */ ]] && SOURCE_TARGET+="/"

# Set kernel source directory
KERNEL_SRC="${SOURCE_TARGET}kernel/kernel-jammy-src"

# Get current running kernel version
RUNNING_KERNEL=$(uname -r)

# Check for kernel source directory
if [ ! -d "$KERNEL_SRC" ]; then
    echo_error "Kernel source not found at $KERNEL_SRC"
    exit 1
fi

# Check for .config
if [ ! -f "$KERNEL_SRC/.config" ]; then
    echo_error "No .config found in $KERNEL_SRC"
    exit 1
fi

# Check for huawei_cdc_ncm.c file
if [ ! -f "$KERNEL_SRC/drivers/net/usb/huawei_cdc_ncm.c" ]; then
    echo_error "huawei_cdc_ncm.c file not found in $KERNEL_SRC/drivers/net/usb/"
    exit 1
fi

# Capture the original directory
ORIGINAL_DIR=$(pwd)

# Set up logging
LOGS_DIR="$ORIGINAL_DIR/logs"
mkdir -p "$LOGS_DIR"
LOG_FILE="$LOGS_DIR/huawei_cdc_ncm_build.log"

# Refresh sudo timestamp
sudo -v

# Change to kernel source directory
cd "$KERNEL_SRC" || {
    echo_error "Failed to change to $KERNEL_SRC"
    exit 1
}

# Check for cdc-wdm dependency
echo_info "Checking for cdc-wdm module dependency..."
if ! grep -q "CONFIG_USB_WDM=\(y\|m\)" .config 2>/dev/null; then
    echo_warning "CONFIG_USB_WDM not enabled in kernel config. Enabling it..."
    sed -i 's/# CONFIG_USB_WDM is not set/CONFIG_USB_WDM=m/' .config 2>/dev/null
    echo "CONFIG_USB_WDM=m" >> .config
    
    echo_info "Compiling cdc-wdm module..."
    sudo make M=drivers/usb/class CONFIG_USB_WDM=m modules 2>&1 | tee -a "$LOG_FILE"
    
    if [ -f "drivers/usb/class/cdc-wdm.ko" ]; then
        echo_success "cdc-wdm module compiled successfully"
    else
        echo_warning "Failed to compile cdc-wdm module. huawei_cdc_ncm may not load properly."
    fi
fi

# Build huawei_cdc_ncm module
echo_info "Building huawei_cdc_ncm module in $KERNEL_SRC"
echo "Building huawei_cdc_ncm module..." | tee -a "$LOG_FILE"

# Determine number of jobs for parallel building
NUM_CPU=$(nproc)
JOBS=$((NUM_CPU > 1 ? NUM_CPU - 1 : 1))

# Try parallel build
time sudo make -j$JOBS M=drivers/net/usb CONFIG_USB_NET_HUAWEI_CDC_NCM=m modules 2>&1 | tee -a "$LOG_FILE"
if [ ${PIPESTATUS[0]} -ne 0 ]; then
    echo_warning "Parallel build failed. Retrying with single-threaded build..." | tee -a "$LOG_FILE"
    time sudo make M=drivers/net/usb CONFIG_USB_NET_HUAWEI_CDC_NCM=m modules 2>&1 | tee -a "$LOG_FILE"
    if [ ${PIPESTATUS[0]} -ne 0 ]; then
        echo_error "Build failed again. Check $LOG_FILE for details"
        exit 1
    fi
fi

# Check if module was built successfully
if [ -f "drivers/net/usb/huawei_cdc_ncm.ko" ]; then
    echo_success "huawei_cdc_ncm module built successfully"
    
    # Prompt for installation
    read -p "Do you want to install the huawei_cdc_ncm module? (y/n): " choice
    if [ "$choice" = "y" ] || [ "$choice" = "Y" ]; then
        echo_info "Installing modules..." | tee -a "$LOG_FILE"
        
        # Install cdc-wdm module if it was built
        if [ -f "drivers/usb/class/cdc-wdm.ko" ]; then
            echo_info "Installing cdc-wdm module..."
            sudo mkdir -p /lib/modules/$RUNNING_KERNEL/kernel/drivers/usb/class/
            sudo cp drivers/usb/class/cdc-wdm.ko /lib/modules/$RUNNING_KERNEL/kernel/drivers/usb/class/
        fi
        
        # Install huawei_cdc_ncm module
        sudo mkdir -p /lib/modules/$RUNNING_KERNEL/kernel/drivers/net/usb/
        sudo cp drivers/net/usb/huawei_cdc_ncm.ko /lib/modules/$RUNNING_KERNEL/kernel/drivers/net/usb/
        sudo depmod -a 2>&1 | tee -a "$LOG_FILE"
        
        echo_success "Modules installed successfully"
        echo_info "You can load the modules with:"
        echo "sudo modprobe cdc-wdm  # Load dependency first"
        echo "sudo modprobe huawei_cdc_ncm"
    else
        echo_info "Skipping installation. To install manually:"
        echo "sudo mkdir -p /lib/modules/$RUNNING_KERNEL/kernel/drivers/usb/class/"
        echo "sudo cp $KERNEL_SRC/drivers/usb/class/cdc-wdm.ko /lib/modules/$RUNNING_KERNEL/kernel/drivers/usb/class/"
        echo "sudo mkdir -p /lib/modules/$RUNNING_KERNEL/kernel/drivers/net/usb/"
        echo "sudo cp $KERNEL_SRC/drivers/net/usb/huawei_cdc_ncm.ko /lib/modules/$RUNNING_KERNEL/kernel/drivers/net/usb/"
        echo "sudo depmod -a"
    fi
else
    echo_error "Failed to build huawei_cdc_ncm module"
    exit 1
fi

# Final log file notification
echo_info "Logs are saved in: $LOG_FILE"

exit 0